package cn.edu.spark.sql

import org.apache.spark.sql.{Encoder, Encoders, SparkSession}
import org.apache.spark.sql.expressions.Aggregator

import java.lang.Thread.sleep

object UserDefinedAggregator2 {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession
      .builder()
      .master("local[*]")
      .appName("Spark SQL user-defined Datasets aggregation example")
      .getOrCreate()

    import spark.implicits._


    val ds = spark.read.json("spark-sql-examples/src/main/resources/employees.json").as[Employee]
    ds.show()
    // +-------+------+
    // |   name|salary|
    // +-------+------+
    // |Michael|  3000|
    // |   Andy|  4500|
    // | Justin|  3500|
    // |  Berta|  4000|
    // +-------+------+

    // Convert the function to a `TypedColumn` and give it a name
    val averageSalary = MyAverage.toColumn.name("average_salary")
    val result = ds.select(averageSalary)
    result.show()
    // +--------------+
    // |average_salary|
    // +--------------+
    // |        3750.0|
    // +--------------+


    sleep(300000)
    spark.stop()
  }


  // typed_custom_aggregation
  case class Employee(name: String, salary: Long)
  case class Average(var sum: Long, var count: Long)


  object MyAverage extends Aggregator[Employee, Average, Double] {
    override def zero: Average = Average(0L, 0L)

    override def reduce(b: Average, a: Employee): Average = {
      b.count += 1L
      b.sum += a.salary
      b
    }

    override def merge(b1: Average, b2: Average): Average = {
      b1.count += b2.count
      b1.sum += b2.sum
      b1
    }

    override def finish(reduction: Average): Double = reduction.sum.toDouble / reduction.count


    override def bufferEncoder: Encoder[Average] = Encoders.product
    override def outputEncoder: Encoder[Double] = Encoders.scalaDouble
  }

}
